package com.selftcp;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;

//TCP服务器
public class TCPServer {
	private final int port;
    private final DatagramSocket socket;
    // 半连接状态的连接
    private final Map<ConnectionKey, TCPConnection> connections = new ConcurrentHashMap<>();
    private volatile boolean running = true;
    // 全连接状态的连接
    private final LinkedBlockingQueue<TCPConnection> newConnections = new LinkedBlockingQueue<>();
    
    public TCPServer(int port) throws IOException {
        this.port = port;
        this.socket = new DatagramSocket(port);
        startListener();
        System.out.println("TCP Server started on port " + port);
    }
    
    // 唯一的Socket读取线程，负责所有数据的接收和分发
    private void startListener() {
        new Thread(() -> {
            try {
                while (running) {
                    byte[] buf = new byte[4096];
                    DatagramPacket packet = new DatagramPacket(buf, buf.length);
                    socket.receive(packet);
                    
                    InetAddress clientAddr = packet.getAddress();
                    int clientPort = packet.getPort();
                    ConnectionKey key = new ConnectionKey(clientAddr, clientPort);
                    
                    TCPacket tcpPacket = TCPacket.fromBytes(packet.getData());
                    
                    if (tcpPacket.syn && !tcpPacket.ackFlag) {
                        // 新的连接请求
                        TCPConnection connection = new TCPConnection(port, clientAddr, clientPort, socket);
                        connections.put(key, connection);
                        connection.handleConnectionRequest(tcpPacket);
                        System.out.println("New connection from " + clientAddr + ":" + clientPort);
                    } else if (connections.containsKey(key)) {
                        // 已存在的连接，分发数据包
                        TCPConnection connection = connections.get(key);
                        if (!tcpPacket.syn && tcpPacket.ackFlag && !connection.isConnected()) {
                            // 完成三次握手
                            if (connection.completeConnection(tcpPacket)) {
                                try {
									newConnections.put(connection);
								} catch (InterruptedException e) {
									e.printStackTrace();
								}
                                System.out.println("Connection established with " + clientAddr + ":" + clientPort);
                            }
                        } else {
                            // 分发数据包到连接
                            connection.dispatchPacket(tcpPacket);
                        }
                    }
                }
            } catch (IOException e) {
                if (running) {
                    e.printStackTrace();
                }
            }
        }).start();
    }
    
    // 接受新连接
    public TCPConnection accept() throws InterruptedException {
        return newConnections.take();
    }
    
    public void stop() {
        running = false;
        socket.close();
        connections.values().forEach(conn -> {
            try {
                conn.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        });
    }
}